import sys
import os
# sys.path.insert(0,'F:\lab\\adv_mnist\DeepRobust-master') # 将该路径插入到包的搜索路径中

sys.path.insert(0, os.path.join(os.getcwd(),)) # 将该路径插入到包的搜索路径中
os.getcwd()

print(sys.path)

import deeprobust.image.netmodels.train_model as trainmodel
'''
说明:若出现“RuntimeError: Given groups=1, weight of size [24, 3, 3, 3], 
expected input[128, 1, 28, 28] to have 3 channels, but got 1 channels instead”
则在对应模型的__init__()方法中将conv1的3改为1,即
self.conv1 = nn.Conv2d(3, num_planes, kernel_size=3, padding=1, bias=False)改为
self.conv1 = nn.Conv2d(1, num_planes, kernel_size=3, padding=1, bias=False)
'''
# trainmodel.train('densenet', 'CIFAR10', 'cuda', 20)
trainmodel.train('ResNet18', 'CIFAR10', 'cuda', 1)
# trainmodel.train('ResNet50', 'CIFAR10', 'cuda', 40)

print('=== end ===')